package org.usmile.algorithms.learn.heap;

import java.util.ArrayList;
import java.util.List;

public class MaxHeap<E extends Comparable<E>> {
    List<E> elems = new ArrayList<>();


    public void add(E e) {
        elems.add(e);

        shiftUp(elems.size() - 1);
    }

    public E getMax() {
        if (size() == 0) {
            return null;
        }

        return elems.get(0);
    }

    public E removeMax() {
        if (size() == 0) {
            return null;
        }

        E max = getMax();

        if (size() == 1) {
            elems.remove(0);

            return max;
        }

        E last = elems.get(size() - 1);
        elems.set(0, last);
        elems.remove(size() - 1);
        shiftDown(0);

        return max;
    }

    private void shiftUp(int index) {
        E e = elems.get(index);
        while (index > 0) {
            int parentIndex = parentIndex(index);
            if (elems.get(parentIndex).compareTo(e) > 0) {
                break;
            } else {
                E parent = elems.get(parentIndex);
                elems.set(index, parent);
            }
            index = parentIndex(index);
        }

        elems.set(index, e);
    }

    private void shiftDown(int index) {
        E e = elems.get(index);
        while (leftChildIndex(index) < size()) {
            int maxChildIndex = leftChildIndex(index);
            if (rightChildIndex(index) < size()) {
                if (elems.get(leftChildIndex(index)).compareTo(elems.get(rightChildIndex(index))) < 0) {
                    maxChildIndex = rightChildIndex(index);
                }
            }
            if (e.compareTo(elems.get(maxChildIndex)) > 0) {
                break;
            }

            elems.set(index, elems.get(maxChildIndex));
            index = maxChildIndex;
        }

        elems.set(index, e);
    }


    public int size() {
        return elems.size();
    }

    private int leftChildIndex(int parentIndex) {
        return 2 * parentIndex + 1;
    }

    private int rightChildIndex(int parentIndex) {
        return 2 * parentIndex + 2;
    }

    private int parentIndex(int childIndex) {
        return (childIndex - 1) / 2;
    }

    public static void main(String[] args) {
        MaxHeap<Integer> maxHeap = new MaxHeap<>();
        maxHeap.add(1);
        maxHeap.add(2);
        maxHeap.add(3);
        maxHeap.add(4);
        maxHeap.add(5);
        System.out.println(maxHeap.removeMax());
        System.out.println(maxHeap.removeMax());
        System.out.println(maxHeap.removeMax());
        System.out.println(maxHeap.removeMax());
        System.out.println(maxHeap.removeMax());
        System.out.println(maxHeap.removeMax());

    }
}
